package leetcode.code1396;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class UndergroundSystem2 {

	Map<Integer, Integer> checkIn1;
	Map<Integer, String> checkIn2;
	Map<String, Map<String, List<Integer>>> records;

	public UndergroundSystem2() {
		checkIn1 = new HashMap<>();
		checkIn2 = new HashMap<>();
		records = new HashMap<>();
	}

	public void checkIn(int id, String stationName, int t) {
		checkIn1.put(id, t);
		checkIn2.put(id, stationName);
	}

	public void checkOut(int id, String stationName, int t) {
		int intime = checkIn1.get(id);
		String intStation = checkIn2.get(id);
		records.computeIfAbsent(intStation, key -> new HashMap<>()).computeIfAbsent(stationName, k -> new ArrayList<>())
				.add(t - intime);
	}

	public double getAverageTime(String startStation, String endStation) {
		Map<String, List<Integer>> m = records.get(startStation);
		List<Integer> list = m.get(endStation);
		double sum = 0d;
		for (double d : list) {
			sum += d;
		}
		return sum / list.size();
	}
}
